import cv2
import numpy as np
from src.data.structures import BinocularImagePair


def visualizeImage(image: np.ndarray, window_name: str = "Image", delay: int = 2000):
    """显示单幅图像

    Args:
        image: 待显示的图像
        window_name: 窗口名称, 默认为 "Image"
        delay: 显示时间 (毫秒), 默认为 2000 毫秒. 设置为 0 则无限等待按键.
    """
    cv2.imshow(window_name, image)
    cv2.waitKey(delay)
    cv2.destroyAllWindows()


def visualizeRectification(
    original_pair: BinocularImagePair,
    rectified_pair: BinocularImagePair,
    num_lines: int = 20,
    delay: int = 2000,
):
    """可视化立体校正前后的效果, 并绘制水平线辅助观察校正效果.

    Args:
        original_pair: 原始的立体图像对
        rectified_pair: 校正后的立体图像对
        num_lines: 要绘制的水平辅助线数量, 默认为20
        delay: 显示时间 (毫秒), 默认为2000毫秒. 设置为0则无限等待按键.

    Note:
        函数会自动处理校正后图像尺寸可能发生变化的情况, 确保显示效果合理
    """
    # 获取原始和校正后图像的尺寸
    h_orig_l, w_orig_l = original_pair.left_image.shape[:2]
    h_orig_r, w_orig_r = original_pair.right_image.shape[:2]
    h_rect_l, w_rect_l = rectified_pair.left_image.shape[:2]
    h_rect_r, w_rect_r = rectified_pair.right_image.shape[:2]

    # 计算显示用的统一尺寸
    w_max = max(w_orig_l + w_orig_r, w_rect_l + w_rect_r)
    h_max = max(h_orig_l, h_orig_r, h_rect_l, h_rect_r)

    # 创建画布
    original_vis = np.zeros((h_max, w_max, 3), dtype=np.uint8)
    rectified_vis = np.zeros((h_max, w_max, 3), dtype=np.uint8)

    # 将原始图像对放入画布
    original_vis[:h_orig_l, :w_orig_l] = original_pair.left_image
    original_vis[:h_orig_r, w_orig_l : w_orig_l + w_orig_r] = original_pair.right_image

    # 将校正后的图像对放入画布, 保持与原始图像相同的布局
    rectified_vis[:h_rect_l, :w_rect_l] = rectified_pair.left_image
    rectified_vis[:h_rect_r, w_orig_l : w_orig_l + w_rect_r] = (
        rectified_pair.right_image
    )  # 使用原始左图宽度作为右图起始位置

    # 在原始图像上画水平线
    for i in range(num_lines):
        y = int(h_max * (i + 1) / (num_lines + 1))
        cv2.line(original_vis, (0, y), (w_max, y), (0, 255, 0), 1)

    # 在校正后的图像上画水平线
    for i in range(num_lines):
        y = int(h_max * (i + 1) / (num_lines + 1))
        cv2.line(rectified_vis, (0, y), (w_max, y), (0, 255, 0), 1)

    # 垂直堆叠显示对比效果
    comparison = np.vstack((original_vis, rectified_vis))

    # 添加文字标注
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(
        comparison, "Original Stereo Pair", (10, 30), font, 1, (220, 120, 70), 2
    )
    cv2.putText(
        comparison,
        "Rectified Stereo Pair",
        (10, h_max + 30),
        font,
        1,
        (220, 120, 70),
        2,
    )

    # 显示结果
    cv2.imshow("Stereo Rectification Comparison", comparison)
    cv2.waitKey(delay)
    cv2.destroyAllWindows()


def visualizeDisparity(
    disparity_image: np.ndarray,
    min_disp: int,
    num_disp: int,
    window_name: str = "Disparity",
    disp_in_rgb: bool = True,
    delay: int = 2000,
):
    """显示单幅视差图像

    Args:
        disparity_image: 待显示的视差图
        min_disp: 最小视差
        num_disp: 视差范围
        window_name: 窗口名称, 默认为 "Image"
        delay: 显示时间 (毫秒), 默认为 2000 毫秒. 设置为 0 则无限等待按键.
    """
    disparity_image = np.nan_to_num(disparity_image, nan=0.0)
    img = (disparity_image - min_disp) / num_disp
    img = np.clip(img, 0, 1)
    img = (img * 255).astype(np.uint8)

    if disp_in_rgb:
        img = cv2.applyColorMap(img, cv2.COLORMAP_JET)

    cv2.imshow(window_name, img)
    cv2.waitKey(delay)
    cv2.destroyAllWindows()
